import numpy as np

from net.LeNet import LeNet5, LeNet5_padding
from net.utils import *
from dataset.preprocess import *
from pipecoco.common import *

from mindspore.nn import SoftmaxCrossEntropyWithLogits
from mindspore.nn import Accuracy
from pipecoco.PipeCoCo import PipeCoCo
from mindspore import context

context.set_context(mode=context.GRAPH_MODE, device_target='CPU')

if __name__=='__main__':

    mnist_path = "../datasets/MNIST_Data"
    model_path = "./models/lenet_padding/"
    network = LeNet5_padding()

    # model_path = "./models/lenet/"
    # network = LeNet5()

    load_model(network, model_path+"/checkpoint_lenet-5_1875.ckpt")


    net_loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')

    ds_eval = create_dataset(os.path.join(mnist_path, "test"),batch_size=32)
    # eval(network, model_path+"/checkpoint_lenet-5_1875.ckpt", net_loss, ds_eval)

    layers_list = []
    dfs_search_layers(network, layers_list)

    ds_test = ds_eval.create_dict_iterator()
    data = next(ds_test)
    input = data["image"].asnumpy()

    # 用pipecoco C-S模型进行计算
    pipecoco_client = PipeCoCo(layers_list, [1 , 32, 32], 2, [0, 4], 'client')
    pipecoco_server = PipeCoCo(layers_list, [1, 32, 32], 2, [0, 4], 'server')

    session_id, init_area_map = pipecoco_client.create_session(input.shape[0])
    pipecoco_server.create_session(input.shape[0], session_id)
    for i in range(4):
        area = init_area_map[i]
        x = input[:,:,area[0][0]:area[0][1],area[1][0]:area[1][1]]
        x = pipecoco_client.fused_layers_forward(x, i)
        pipecoco_server.fused_layers_forward(x, i)

    x = pipecoco_server.non_fused_layers_forward()
    pred = np.argmax(x, axis=1)
    print(pred)

    load_model(network, model_path + "/checkpoint_lenet-5_1875.ckpt")
    model = Model(network, net_loss, metrics={"Accuracy": Accuracy()})
    label = np.argmax(model.predict(data["image"]), axis=1)
    print(label)